# from numpy.distutils.cpuinfo import cpu
from torch import Tensor, nn
from CNN import cnnNet
from torch.autograd import Variable
import torch
import torch.utils.data as Data
import numpy as np
from MyDataset import MyDataset
import matplotlib.pyplot as plt

real_epochs = 0
one_one = 0
one_zero = 0
zero_zero = 0
zero_one = 0
# 实例化
cnnModel = cnnNet()
epochs = 250
# batch_size = 128
batch_size = 512
lr = 0.0001
# loss_list = np.array([])
loss_list = []
best_loss = 100
es = 0

# 定义损失函数和
loss_func = nn.BCELoss(reduction='mean')
optimizer = torch.optim.Adam(cnnModel.parameters(), lr=lr, weight_decay=0.00001)

torch.set_grad_enabled(True)
cnnModel.train()
# device = torch.device("cuda:0")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cnnModel.to(device)

# 获取数据集
myTrainDataset = MyDataset("D:\PycharmProjects\My_project\kneepoint\csv\\train", "D:\PycharmProjects\My_project\kneepoint\data\\train")
print(myTrainDataset)
# 读取数据集
train_loader = torch.utils.data.DataLoader(
    dataset=myTrainDataset,
    batch_size=batch_size,
    shuffle=True
)
for epoch in range(epochs):
    for batch_idx, data in enumerate(train_loader):
        # print(batch_idx)
        label, txtData = data
        txtData = txtData.to(device, torch.float)
        optimizer.zero_grad()
        label = label.to(device, torch.float)

        classify_pre = cnnModel(txtData)

        classify_pre = torch.squeeze(classify_pre)

        # print(label)
        # print(classify_pre)
        loss = loss_func(classify_pre, label)
        loss_list.append(Tensor.cpu(loss).item())
        print("list: ", loss_list)

        print(epoch, 'loss:' + str(loss))
        loss.backward()
        optimizer.step()

x = np.linspace(1, epochs, epochs)
# print(x)
# print(type(x))
# x = list(x)

# loss_list = loss_list.detach().numpy()

print(x.shape, len(loss_list))
print(loss_list)

plt.plot(x, loss_list)
plt.ylabel('loss')
plt.xlabel('epoch')
plt.grid()
plt.savefig('image/loss/' + str(epoch) + '_loss.png')
plt.show()
# torch.save(cnnModel, './LeNet.pkl')

# print(loss_list)
# fx2 = sympy.integrate((((y-r2)/(l2*l2))*sympy.exp(-(y-r2)*(y-r2)/(2*l2*l2))), (y, 0, m))-b
# w = sympy.solve(fx2, m)
# print(w)
# z=bf2(c,r3,l3,k,-2e-2,0)
# print(z)
